package com.freestyle.sentinel;

import com.alibaba.csp.sentinel.Entry;
import com.alibaba.csp.sentinel.EntryType;
import com.alibaba.csp.sentinel.SphU;
import com.alibaba.csp.sentinel.annotation.SentinelResource;
import com.alibaba.csp.sentinel.annotation.aspectj.SentinelResourceAspect;
import com.alibaba.csp.sentinel.slots.block.BlockException;
import com.freestyle.sentinel.annotation.SentinelResourceGlobalHandle;
import com.freestyle.sentinel.com.freestyle.sentinel.GlobalHandler;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Pointcut;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * Created by rocklee on 2022/3/28 16:10
 */
public class SentinelResourceAspectEx extends SentinelResourceAspect {
  private ConcurrentHashMap<Method,Method> blockHandler=new ConcurrentHashMap<>();
  private ConcurrentHashMap<Method,Method> fallbackHandler=new ConcurrentHashMap<>();
  private GlobalHandler globalHandler;


  public SentinelResourceAspectEx(GlobalHandler globalHandler) {
    this.globalHandler = globalHandler;
  }

  private Method getBlockHandler(Method orgMethod, SentinelResourceGlobalHandle annotation){
    if (annotation.auto()){
        if (globalHandler==null){
          throw new IllegalStateException("globalHandler is null , can't set SentinelResourceEx.auto equals true");
        }
        return  blockHandler.computeIfAbsent(orgMethod,(method)->
        {
          try {
            return globalHandler.getClass().getMethod("globalBlockHandler", BlockException.class,MethodArgs.class);
          } catch (NoSuchMethodException e) {
            e.printStackTrace();
          }
          return null;
        }
        );
    }
    return null;
  }
  private Method getFallBackHandler(Method orgMethod, SentinelResourceGlobalHandle annotation){
    if (annotation.auto()){
      if (globalHandler==null){
        throw new IllegalStateException("globalHandler is null , can't set SentinelResourceEx.auto equals true");
      }
      return fallbackHandler.computeIfAbsent(orgMethod,method -> {
        try {
          return globalHandler.getClass().getMethod("globalFallBackHandler", Throwable.class, MethodArgs.class);
        } catch (NoSuchMethodException e) {
        }
        return null;
      });
    }
    return null;
  }
  private boolean procGlobalHandle(Method originMethod, ProceedingJoinPoint pjp, Throwable ex,boolean fallbackOnly, Consumer<Object> onReturn) throws InvocationTargetException, IllegalAccessException {
    SentinelResourceGlobalHandle sentinelResourceEx=originMethod.getAnnotation(SentinelResourceGlobalHandle.class);
    if (sentinelResourceEx!=null) {
      Method blockMethod = getBlockHandler(originMethod, sentinelResourceEx);
      if (blockMethod!=null&&!fallbackOnly){
        onReturn.accept(blockMethod.invoke(globalHandler,ex,new MethodArgs(originMethod,pjp.getArgs(),originMethod.getGenericReturnType())));
        return true;
      }
      Method fallbackMethod=getFallBackHandler(originMethod,sentinelResourceEx);
      if (fallbackMethod!=null ){
        onReturn.accept(fallbackMethod.invoke(globalHandler,ex,new MethodArgs(originMethod,pjp.getArgs(),originMethod.getGenericReturnType())));
        return true;
      }
    }
    return false;
  }
  @Pointcut(value = "@annotation(com.freestyle.sentinel.annotation.SentinelResourceGlobalHandle)")
  @Override
  public void sentinelResourceAnnotationPointcut() {
  }
  @Around("sentinelResourceAnnotationPointcut()")
  public Object invokeResourceWithSentinel(ProceedingJoinPoint pjp) throws Throwable {
    Method originMethod = resolveMethod(pjp);

    SentinelResource annotation = originMethod.getAnnotation(SentinelResource.class);
    if (annotation == null) {
      // Should not go through here.
      throw new IllegalStateException("Wrong state for SentinelResource annotation");
    }
    String resourceName = getResourceName(annotation.value(), originMethod);
    EntryType entryType = annotation.entryType();
    int resourceType = annotation.resourceType();
    Entry entry = null;
    try {
      entry = SphU.entry(resourceName, resourceType, entryType, pjp.getArgs());
      Object result = pjp.proceed();
      return result;
    } catch (BlockException ex) {
      final Object[] val = new Object[1];
      if (procGlobalHandle(originMethod,pjp,ex,false,(obj)->{
        val[0] =obj;
      })){
        return val[0];
      };
      return handleBlockException(pjp, annotation, ex);
    } catch (Throwable ex) {
      Class<? extends Throwable>[] exceptionsToIgnore = annotation.exceptionsToIgnore();
      // The ignore list will be checked first.
      if (exceptionsToIgnore.length > 0 && exceptionBelongsTo(ex, exceptionsToIgnore)) {
        throw ex;
      }
      if (exceptionBelongsTo(ex, annotation.exceptionsToTrace())) {
        traceException(ex);
        final Object[] val = new Object[1];
        if (procGlobalHandle(originMethod,pjp,ex,true,(obj)->{
          val[0] =obj;
        })){
          return val[0];
        };
        return handleFallback(pjp, annotation, ex);
      }

      // No fallback function can handle the exception, so throw it out.
      throw ex;
    } finally {
      if (entry != null) {
        entry.exit(1, pjp.getArgs());
      }
    }
  }
}
